import operator
import sqlite3

from typing import Annotated,Sequence

from langgraph.types import RetryPolicy
from typing_extensions import TypedDict

from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage

from langgraph.graph import START,StateGraph,END


#定义数据库和模型
db = SQLDatabase.from_uri("sqlite:///:memory:")
model = ChatOpenAI(
        model="qwen-plus",  # 阿里云千问-plus模型
        openai_api_key='sk-965dc39b016c49ecbe29de180f4db2b6',
        openai_api_base='https://dashscope.aliyuncs.com/compatible-mode/v1',
        temperature=0.7,  # 控制生成文本的随机性
        max_tokens=2048,  # 最大生成长度
    )
#定义图的状态和逻辑节点
class AgentState(TypedDict):
    messages:Annotated[Sequence[BaseMessage],operator.add]

def query_database(state):
    query_result = db.run("select * from Artist limit 10;")
    return {"messages":[AIMessage(content=query_result)]}

def call_model(state):
    user_input = str(state['messages'][0]['content'])
    # 正确的消息格式
    messages = [
        SystemMessage(content="You are a helpful assistant."),
        HumanMessage(content="Hello, how are you?")
    ]
    response = model.invoke(messages)
    return {"messages":[response]}

#定义图
builder = StateGraph(AgentState)

builder.add_node("query_database",
                 query_database,
                 re=RetryPolicy(retry_on=sqlite3.OperationalError)
                 )
builder.add_node("model",call_model,
                 retry=RetryPolicy(max_attempts=5)
                 )
builder.add_edge(START,"model")
builder.add_edge("model","query_database")
builder.add_edge("query_database",END)
graph = builder.compile()
#调用图
initial_state = {"messages": [{"role": "user", "content": ["你好"]}]}
result = graph.invoke(initial_state)
